import os
import re
import threading
import queue
import requests
import m3u8
from lxml import etree
from urllib.parse import urljoin, urlparse
from Crypto.Cipher import AES
from tkinter import *
from tkinter import ttk, scrolledtext, messagebox
import subprocess

m3u8_queue = queue.Queue()

def get_html_tree(url):
    resp = requests.get(url, timeout=10)
    resp.encoding = resp.apparent_encoding
    return etree.HTML(resp.text)

def crawl_video_detail_urls(start_url, video_xpath, next_page_xpath):
    video_detail_urls = []
    url = start_url
    while url:
        print(f"[列表页] 抓取: {url}")
        html = get_html_tree(url)
        video_links = html.xpath(video_xpath)
        for link in video_links:
            full_link = urljoin(url, link)
            video_detail_urls.append(full_link)
            print(f"  视频详情页: {full_link}")
        next_links = html.xpath(next_page_xpath)
        if next_links:
            url = urljoin(url, next_links[0])
        else:
            url = None
    return video_detail_urls

def extract_m3u8_from_detail(detail_url, m3u8_regex=r'(https?://[^\s\'"]+\.m3u8)'):
    try:
        resp = requests.get(detail_url, timeout=10)
        text = resp.text
        m3u8_links = re.findall(m3u8_regex, text)
        print(f"[详情页] {detail_url} -> 找到 {len(m3u8_links)} 个 m3u8")
        return m3u8_links
    except Exception as e:
        print(f"[详情页] {detail_url} 抓取失败: {e}")
        return []

class M3U8DownloaderGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("🐍 M3U8 多级抓取下载器")
        self.root.geometry("900x700")

        # 主容器，使用顶层Frame管理上下结构
        main_frame = Frame(root)
        main_frame.pack(fill=BOTH, expand=True)

        # 上半部分：输入区，固定高度，放单选和输入框
        input_frame = Frame(main_frame)
        input_frame.pack(side=TOP, fill=X, padx=10, pady=10)

        # 下载方式选择
        self.mode_var = IntVar(value=1)
        frame_mode = Frame(input_frame)
        frame_mode.pack(anchor=W)
        Radiobutton(frame_mode, text="通过视频列表页下载", variable=self.mode_var, value=1, command=self.toggle_mode).pack(side=LEFT, padx=10)
        Radiobutton(frame_mode, text="直接输入 m3u8 链接下载", variable=self.mode_var, value=2, command=self.toggle_mode).pack(side=LEFT, padx=10)

        # 视频列表页相关输入
        self.frame_list = Frame(input_frame)
        self.frame_list.pack(anchor=W, fill=X, pady=5)

        Label(self.frame_list, text="视频列表页 URL:").pack(anchor=W)
        self.url_entry = Entry(self.frame_list, width=110)
        self.url_entry.pack(pady=2)

        Label(self.frame_list, text="视频详情页 XPath (提取视频链接):").pack(anchor=W)
        self.video_xpath_entry = Entry(self.frame_list, width=110)
        self.video_xpath_entry.pack(pady=2)

        Label(self.frame_list, text="下一页 XPath (翻页链接):").pack(anchor=W)
        self.next_xpath_entry = Entry(self.frame_list, width=110)
        self.next_xpath_entry.pack(pady=2)

        # m3u8 直链输入（初始隐藏）
        self.frame_m3u8 = Frame(input_frame)
        Label(self.frame_m3u8, text="直接输入 m3u8 链接:").pack(anchor=W)
        self.m3u8_entry = Entry(self.frame_m3u8, width=110)
        self.m3u8_entry.pack(pady=2)
        self.frame_m3u8.pack_forget()

        # 开始按钮，放在输入区底部
        self.start_button = Button(input_frame, text="开始抓取并下载", command=self.start_download)
        self.start_button.pack(pady=10)

        # 下半部分：进度条和日志，占满剩余空间，且日志可伸缩
        output_frame = Frame(main_frame)
        output_frame.pack(side=TOP, fill=BOTH, expand=True, padx=10, pady=5)

        self.progress = ttk.Progressbar(output_frame, orient="horizontal", length=850, mode="determinate")
        self.progress.pack(fill=X, pady=5)

        Label(output_frame, text="日志:").pack(anchor=W)
        self.log = scrolledtext.ScrolledText(output_frame, width=110, height=30)
        self.log.pack(fill=BOTH, expand=True)

        # 文件夹准备
        self.save_dir = "downloads"
        self.output_dir = "segments"
        os.makedirs(self.save_dir, exist_ok=True)
        os.makedirs(self.output_dir, exist_ok=True)

    def toggle_mode(self):
        mode = self.mode_var.get()
        if mode == 1:
            self.frame_list.pack(anchor=W, fill=X, padx=10)
            self.frame_m3u8.pack_forget()
        else:
            self.frame_list.pack_forget()
            self.frame_m3u8.pack(anchor=W, fill=X, padx=10)

    def log_msg(self, msg):
        self.log.insert(END, msg + "\n")
        self.log.see(END)

    def start_download(self):
        mode = self.mode_var.get()
        if mode == 1:
            list_url = self.url_entry.get().strip()
            video_xpath = self.video_xpath_entry.get().strip()
            next_xpath = self.next_xpath_entry.get().strip()
            if not list_url or not video_xpath:
                messagebox.showerror("错误", "请填写列表页 URL 和视频详情页 XPath")
                return
            self.log_msg(f"开始抓取视频列表页: {list_url}")
            threading.Thread(target=self.run_crawl_and_download,
                             args=(list_url, video_xpath, next_xpath), daemon=True).start()
        else:
            m3u8_url = self.m3u8_entry.get().strip()
            if not m3u8_url or not m3u8_url.endswith(".m3u8"):
                messagebox.showerror("错误", "请输入正确的 m3u8 链接")
                return
            self.log_msg(f"直接下载 m3u8: {m3u8_url}")
            threading.Thread(target=self.run_direct_m3u8_download,
                             args=(m3u8_url,), daemon=True).start()

    def run_crawl_and_download(self, list_url, video_xpath, next_xpath):
        detail_urls = crawl_video_detail_urls(list_url, video_xpath, next_xpath)
        self.log_msg(f"共抓取到 {len(detail_urls)} 个视频详情页链接")
        all_m3u8_links = []
        for idx, detail_url in enumerate(detail_urls, 1):
            m3u8s = extract_m3u8_from_detail(detail_url)
            all_m3u8_links.extend(m3u8s)
            self.log_msg(f"[详情页 {idx}/{len(detail_urls)}] {detail_url} 找到 {len(m3u8s)} 个 m3u8")
        if not all_m3u8_links:
            self.log_msg("未找到任何 m3u8 链接，结束")
            return
        self.log_msg(f"总共抓取到 {len(all_m3u8_links)} 个 m3u8 链接，开始下载")
        self.progress['maximum'] = len(all_m3u8_links)
        self.progress['value'] = 0
        threading.Thread(target=self.producer, args=(all_m3u8_links,), daemon=True).start()
        for _ in range(3):
            threading.Thread(target=self.consumer, daemon=True).start()

    def run_direct_m3u8_download(self, m3u8_url):
        self.progress['maximum'] = 1
        self.progress['value'] = 0
        threading.Thread(target=self.producer, args=([m3u8_url],), daemon=True).start()
        for _ in range(3):
            threading.Thread(target=self.consumer, daemon=True).start()

    def producer(self, urls):
        for url in urls:
            try:
                resp = requests.get(url, timeout=15)
                resp.raise_for_status()
                filename = os.path.basename(urlparse(url).path)
                if not filename.endswith('.m3u8'):
                    filename += '.m3u8'
                file_path = os.path.join(self.save_dir, filename)
                with open(file_path, 'w', encoding='utf-8') as f:
                    f.write(resp.text)
                m3u8_queue.put((url, file_path))
                self.log_msg(f"[Producer] 已放入队列: {file_path}")
            except Exception as e:
                self.log_msg(f"[Producer] 下载 m3u8 失败: {e}")

    def consumer(self):
        while True:
            try:
                url, file_path = m3u8_queue.get(timeout=30)
                self.log_msg(f"[Consumer] 处理: {file_path}")
                self.process_m3u8(url, file_path)
                self.progress['value'] += 1
            except queue.Empty:
                self.log_msg("[Consumer] 队列空，下载结束")
                break

    def download_and_decrypt_segment(self, segment_uri, base_uri, key, iv, idx):
        segment_url = urljoin(base_uri, segment_uri)
        try:
            resp = requests.get(segment_url, timeout=15)
            resp.raise_for_status()
            content = resp.content
            if key:
                cipher = AES.new(key, AES.MODE_CBC, iv)
                content = cipher.decrypt(content)

            # 基础文件名
            base_filename = f"{idx:05d}.ts"
            save_path = os.path.join(self.output_dir, base_filename)

            # 如果文件存在，循环改名
            count = 1
            while os.path.exists(save_path):
                name, ext = os.path.splitext(base_filename)
                new_name = f"{name}_{count}{ext}"
                save_path = os.path.join(self.output_dir, new_name)
                count += 1

            with open(save_path, 'wb') as f:
                f.write(content)
        except Exception as e:
            self.log_msg(f"[Segment] 下载失败 {segment_url}: {e}")

    def process_m3u8(self, m3u8_url, m3u8_file):
        try:
            m3u8_obj = m3u8.load(m3u8_file)
        except Exception as e:
            self.log_msg(f"[Process] 解析 m3u8 失败: {e}")
            return

        # base_uri = m3u8_obj.base_uri or os.path.dirname(m3u8_url) + '/'
        # 确保 base_uri 是完整URL目录
        if m3u8_obj.base_uri and m3u8_obj.base_uri.startswith('http'):
            base_uri = m3u8_obj.base_uri
        else:
            # 如果m3u8_obj.base_uri为空，自己从m3u8_url提取目录路径
            parsed = urlparse(m3u8_url)
            base_uri = f"{parsed.scheme}://{parsed.netloc}{os.path.dirname(parsed.path)}/"
        segments = m3u8_obj.segments

        key = None
        iv = None
        if m3u8_obj.keys and m3u8_obj.keys[0]:
            key_uri = m3u8_obj.keys[0].uri
            key_url = urljoin(base_uri, key_uri)
            try:
                self.log_msg(f"[Process] 检测到加密，下载密钥")
                key = requests.get(key_url, timeout=15).content
                iv = bytes.fromhex(m3u8_obj.keys[0].iv[2:]) if m3u8_obj.keys[0].iv else b'\x00' * 16
            except Exception as e:
                self.log_msg(f"[Process] 下载密钥失败: {e}")
                return

        threads = []
        for idx, segment in enumerate(segments):
            t = threading.Thread(target=self.download_and_decrypt_segment,
                                 args=(segment.uri, base_uri, key, iv, idx))
            t.start()
            threads.append(t)
            if len(threads) >= 16:
                for tt in threads:
                    tt.join()
                threads = []
        for tt in threads:
            tt.join()

        merged_file = os.path.join(self.save_dir, f"{os.path.basename(m3u8_file)}.ts")
        self.merge_ts_files(merged_file)

        mp4_file = merged_file.replace('.ts', '.mp4')
        self.ffmpeg_remux(merged_file, mp4_file)
        try:
            os.remove(merged_file)
        except:
            pass
        self.log_msg(f"[Process] 合并完成：{mp4_file}")

    def merge_ts_files(self, output_file):
        ts_files = sorted([f for f in os.listdir(self.output_dir) if f.endswith('.ts')])
        with open(output_file, 'wb') as outfile:
            for ts in ts_files:
                path = os.path.join(self.output_dir, ts)
                with open(path, 'rb') as infile:
                    outfile.write(infile.read())
                try:
                    os.remove(path)
                except:
                    pass

    def ffmpeg_remux(self, input_ts, output_mp4):
        ffmpeg_exe = os.path.join(os.getcwd(), 'ffmpeg.exe')
        if not os.path.isfile(ffmpeg_exe):
            self.log_msg("[ffmpeg] 未找到 ffmpeg.exe，跳过封装")
            return
        cmd = [ffmpeg_exe, '-y', '-i', input_ts, '-c', 'copy', output_mp4]
        try:
            subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        except Exception as e:
            self.log_msg(f"[ffmpeg] 封装失败: {e}")

if __name__ == '__main__':
    # https://test-streams.mux.dev/x36xhzz/url_6/193039199_mp4_h264_aac_hq_7.m3u8
    # https://test-streams.mux.dev/x36xhzz/url_0/193039199_mp4_h264_aac_hd_7.m3u8
    root = Tk()
    app = M3U8DownloaderGUI(root)
    root.mainloop()
